﻿#pragma once
#include "ai_global.h"
#include "IInfer.h"

#include <opencv2/opencv.hpp>
#include <onnxruntime_cxx_api.h>
#include "AIDataDefine.h"
//#include <string>

class AI_EXPORT InferObjDetect :
    public IInfer
{
public:
	InferObjDetect();
	~InferObjDetect();
	int infer(cv::Mat& src);
	InferResult* getInferResult() { return &result_; }
	bool loadModelFile(const char* modelPath);
	bool loadLabelFile(const char* filePath);
	void setParams(float confidence = 0.25, float score = 0.25, float nms = 0.45);
	std::string getLastError() { return lastErrorInfo_; }
	std::vector<std::string> getTestOutput() { return testOutputInfos_; }

protected:
	void release();

	// classesFile
	void parseTextFile(std::string filePath);
	void parseYamlFile(std::string filePath);

private:
	InferResult result_;
	std::string lastErrorInfo_ = "";
	std::vector<std::string> testOutputInfos_;

	// onnx相关
	Ort::Session* session_ = nullptr;
	Ort::SessionOptions* session_options_ = nullptr;;
	Ort::Env* env_ = nullptr;

	float xFactor_ = 1.0;
	float yFactor_ = 1.0;
	int input_w = 0;
	int input_h = 0;
	int output_h = 0;
	int output_w = 0;

	std::vector<std::string> inputNodeNames_;
	std::vector<std::string> outputNodeNames_;

	std::vector<std::string> labels_;

	float confidenceThresh_ = 0.25;
	float scoreThresh_ = 0.25;
	float nmsThresh_ = 0.45;
};

